from django.contrib import admin
from user.models import UserProfile
from django.utils.translation import gettext_lazy
from knowledge1.models import Base, Detail, Project, Source
from text_embedding import embedding_model_dict, choose_model
from collections import defaultdict
from knowledge1.project_file.conn_milvus import MilvusDB
from datetime import datetime
import uuid


class DetailAdmin(admin.ModelAdmin):
    list_display = ('uid', 'base', 'text', 'detail_text', 'source',  'embedding', 'meta', 'parent_id', 'comment', 'create_time',)
    list_filter = ['uid', 'base', 'source', 'embedding', 'comment']
    search_fields = ['text']
    readonly_fields = ['uid', 'parent_id', 'create_time']

    list_per_page = 20
    SIMPLEUI_LOGO = "fa fa-envelope"

    # def has_add_permission(self, request=None):
    #     # Disable add
    #     return False

    # def has_change_permission(self, request, obj=None):
    #     return True

    # def has_delete_permission(self, request, obj=None):
    #     return False
    
    def get_queryset(self, request):
        qs = super(DetailAdmin, self).get_queryset(request)
        if request.user.is_superuser:
            return qs.filter(status=1)

        knowledge_permission = UserProfile.objects.get(username=request.user.username).knowledge_permission.all()
        return qs.filter(base__project__in=knowledge_permission, status=1)
    
    # 外部
    def delete_queryset(self, request, queryset):
        project_dict = defaultdict(list)
        base_dict = defaultdict(list)

        # mysql数据库删除
        # 按照项目做划分
        for obj in queryset:
            project_name = obj.base.project.name
            embedding_dim = obj.embedding.dims
            project_dim_name = project_name + '_' + str(embedding_dim)
            project_dict[project_dim_name].append(obj)
        for project_dim_name, del_obj_list in project_dict.items():
            db_IP_name = project_dim_name + '_IP'
            db_L2_name = project_dim_name + '_L2'
            db_IP = MilvusDB(collection_name=db_IP_name, index_type='IP', dimension=embedding_dim)
            db_L2 = MilvusDB(collection_name=db_L2_name, index_type='L2', dimension=embedding_dim)
            db_list = [db_IP, db_L2]
                
            for del_obj in del_obj_list:
                base_name = del_obj.base.name
                uid = del_obj.uid
                base_dict[base_name].append(uid)
            for base_name, del_uid_list in base_dict.items():
                for vector_db in db_list:
                    return_status = vector_db.delete_data([del_uid_list], partition_name=base_name)
            if return_status:
                for obj in queryset:
                    obj.delete()
                message_bit = "%s 条记录被" % (queryset.count())
                self.message_user(request, "%s成功删除." % message_bit)
            else:
                self.message_user(request, "删除失败.")
    
    # 内部
    def delete_model(self, request, obj):
        # 按照项目做划分
        uid = obj.uid
        project_name = obj.base.project.name
        base_name = obj.base.name
        embedding_dim = obj.embedding.dims
        project_dim_name = project_name + '_' + str(embedding_dim)
        
        db_IP_name = project_dim_name + '_IP'
        db_L2_name = project_dim_name + '_L2'
        db_IP = MilvusDB(collection_name=db_IP_name, index_type='IP', dimension=embedding_dim)
        db_L2 = MilvusDB(collection_name=db_L2_name, index_type='L2', dimension=embedding_dim)
        db_list = [db_IP, db_L2]
                
        for vector_db in db_list:
            return_status = vector_db.delete_data([[uid]], partition_name=base_name)
        if return_status:
            obj.delete()
            message_bit = "1 条记录被"
            self.message_user(request, "%s成功删除." % message_bit)
        else:
            self.message_user(request, "删除失败.")
            
            
    def save_model(self, request, obj, form, change):
        project_name = obj.base.project.name
        embedding_type = obj.embedding.name
        base_name = obj.base.name
        meta_data = str(obj.meta)
        tmp_embedding_dim = obj.embedding.dims
        source = obj.source.name
        project_dim_name = project_name + '_' + str(tmp_embedding_dim)
        text = obj.text

        db_IP_name = project_dim_name + '_IP'
        db_L2_name = project_dim_name + '_L2'
        db_IP = MilvusDB(collection_name=db_IP_name, index_type='IP', dimension=tmp_embedding_dim)
        db_L2 = MilvusDB(collection_name=db_L2_name, index_type='L2', dimension=tmp_embedding_dim)
        db_list = [db_IP, db_L2]
        # 修改数据
        if change:
            # 在milvus中删除数据
            for vector_db in db_list:
                return_status = vector_db.delete_data([[obj.uid]], partition_name=base_name)
                if return_status:
                    print("修改数据")
                    # 写入milvus
                    model = choose_model(embedding_type)
                    query_embedding = model(text)
                    tmp_result_list = [
                        [obj.uid],
                        [text],
                        [source],
                        [meta_data],
                        [query_embedding],
                        [embedding_type],
                        [project_name],
                        [base_name],
                        [datetime.now().strftime("%Y-%m-%d %H:%M:%S")]
                    ]
                    return_status, error_info = vector_db.insert_data(data=tmp_result_list, partition_name=base_name)
                if return_status:
                    obj.save()
                    self.message_user(request, "修改成功.")
        # 新增数据
        else:
            obj.uid = datetime.now().strftime('%Y%m%d%H%M%S') + uuid.uuid4().hex
            for vector_db in db_list:
                print("增加数据")
                # 写入milvus
                model = choose_model(embedding_type)
                query_embedding = model(text)
                tmp_result_list = [
                    [obj.uid],
                    [text],
                    [source],
                    [meta_data],
                    [query_embedding],
                    [embedding_type],
                    [project_name],
                    [base_name],
                    [datetime.now().strftime("%Y-%m-%d %H:%M:%S")]
                ]
                return_status, error_info = vector_db.insert_data(data=tmp_result_list, partition_name=base_name)
            if return_status:
                obj.save()
                self.message_user(request, "新增成功.")
        print("后台直接写入数据！")

    fieldsets = (
        (gettext_lazy('知识详情'), {'fields': ('text', 'detail_text',)}),
        (gettext_lazy('知识其他信息'), {'fields': ('source', 'embedding','meta')}),
        (gettext_lazy('所属项目'), {'fields': ('base',)}),
    )


class ProjectAdmin(admin.ModelAdmin):
    list_display = ('name', 'comment', 'create_time',)
    list_filter = ['name', 'comment']
    search_fields = ['name']
    readonly_fields = ['name', 'create_time']

    list_per_page = 10
    SIMPLEUI_LOGO = "fa fa-envelope-open"

    def has_add_permission(self, request=None):
        # Disable add
        return True

    # def has_change_permission(self, request, obj=None):
    #     return False

    def has_delete_permission(self, request, obj=None):
        return False

    def get_queryset(self, request):
        qs = super(ProjectAdmin, self).get_queryset(request)
        # if request.user.is_superuser:
        #     return qs
        return qs.filter(status=1)

class BaseAdmin(admin.ModelAdmin):
    list_display = ('name',  'project', 'category', 'group_name', 'create_time',)
    list_filter = ['name', 'project']
    search_fields = ['name', 'category', 'project__name']
    readonly_fields = ['name',  'project', 'create_time',]
    list_per_page = 10
    SIMPLEUI_LOGO = "fa fa-envelope-open"

    def has_add_permission(self, request=None):
        # Disable add
        return False

    def has_change_permission(self, request, obj=None):
        return True

    def has_delete_permission(self, request, obj=None):
        return False

    def get_queryset(self, request):
        qs = super(BaseAdmin, self).get_queryset(request)
        # if request.user.is_superuser:
        #     return qs
        return qs.filter(status=1)


class SourceAdmin(admin.ModelAdmin):
    list_display = ('name', 'status', 'create_time',)
    list_filter = ['name', 'status', ]
    search_fields = ['name']
    readonly_fields =  ['create_time']
    list_per_page = 10
    SIMPLEUI_LOGO = "fa fa-envelope-open"

    def has_add_permission(self, request=None):
        # Disable add
        return True

    def has_change_permission(self, request, obj=None):
        return True

    def has_delete_permission(self, request, obj=None):
        return False

    def get_queryset(self, request):
        qs = super(SourceAdmin, self).get_queryset(request)
        # if request.user.is_superuser:
        #     return qs
        return qs.filter(status=1)


class LabelAdmin(admin.ModelAdmin):
    list_display = ('name', 'project', 'source', 'status', 'create_time',)
    list_filter = ['name', 'project', 'source', 'status']
    search_fields = ['name', 'project__name', 'source__name']
    readonly_fields = ['project', 'source','create_time']
    list_per_page = 10
    SIMPLEUI_LOGO = "fa fa-envelope-open"

    def has_add_permission(self, request=None):
        # Disable add
        return True

    def has_change_permission(self, request, obj=None):
        return True

    def has_delete_permission(self, request, obj=None):
        return False

    def get_queryset(self, request):
        qs = super(LabelAdmin, self).get_queryset(request)
        # if request.user.is_superuser:
        #     return qs
        return qs.filter(status=1)


admin.site.register(Detail, DetailAdmin)
admin.site.register(Project, ProjectAdmin)
admin.site.register(Base, BaseAdmin)
admin.site.register(Source, SourceAdmin)
# admin.site.register(Label, LabelAdmin)


